import os

from langchain_community.chat_models import ChatZhipuAI
from openai import OpenAI
import json

# openai_api_base = os.environ['OPENAI_API_BASE']
# openai_api_key = os.environ['OPENAI_API_KEY']
# client = OpenAI(
#     api_key=openai_api_key,
#     base_url=openai_api_base
# )
base_url = 'https://api.lkeap.cloud.tencent.com/v1'
api_key = 'sk-7EIVZor7TFz2e9ZmXEcKacVPEOcgUphCIxfeMIh9vp8Xxr4w'

client = OpenAI(
    api_key=api_key,
    base_url=base_url
)



def weather_function(location):
    match location:
        case "无锡" | "wuxi":
            weather = "晴天，下雪"
        case "苏州" | "suzhou":
            weather = "多云"
        case "常州" | "changzhou":
            weather = "雨"
        case _:
            weather = "不清楚"

    weather_answer = [
        {"天气": weather}
    ]

    return json.dumps(weather_answer)


functions = [
    {
        "type": "function",
        "function": {
            "name": "weather",
            "description": "根据位置了解天气",
            "parameters": {
                "type": "object",
                "properties": {

                    "location": {
                        "type": "string",
                        "description": "输入您想要了解天气的位置。 示例：东京",
                    },
                },
                "required": ["location"],
            }
        }
    }
]


def get_completion(message):
    response = client.chat.completions.create(
        # model="Qwen/Qwen2.5-14B-Instruct",
        model="deepseek-v3",
        messages=message,
        tools=functions,
        function_call="auto",
    )
    message = response.choices[0].message
    print(message)
    return message


messages = [
    {
        "role": "system",
        "content": "You are a useful assistant."
    },
    {
        "role": "user",
        "content": "无锡天气怎么样?"
    },
]
print(messages[1]["content"])

response = get_completion(messages)
messages.append(response)

# 如果LLM需要function calling，调用相应的函数，并将函数结果数据加入对话上下文，继续调用LLM。
while (response.tool_calls is not None):
    for tool_call in response.tool_calls:
        function_name = tool_call.function.name
        arguments = json.loads(tool_call.function.arguments)
        if (function_name == "weather"):
            weatherNow = weather_function(location=arguments.get('location'))
            # messages.append(tool_call)
            # messages.append({"role": "function", "name": "weather", "content": weatherNow})

            messages.append({
                "tool_call_id": tool_call.id,  # 用于标识函数调用的 ID
                "role": "tool",
                "name": tool_call.function.name,
                "content": str(weatherNow)  # 数值result 必须转成字符串
            })
            print(messages)
            response = get_completion(messages)
            messages.append(response)
            print(messages)

